import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.MemoryLayout;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.StructLayout;
import java.lang.foreign.SymbolLookup;
import java.lang.invoke.MethodHandle;
import java.nio.charset.StandardCharsets;
import static java.lang.foreign.ValueLayout.ADDRESS;
import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static java.lang.foreign.ValueLayout.JAVA_INT;
import static java.lang.foreign.ValueLayout.JAVA_LONG;

// javac: --enable-preview -source 21 -J-Duser.language=en
// java : --enable-preview --enable-native-access=ALL-UNNAMED

/**
 * @see <a href="https://openjdk.org/jeps/442">Foreign Function & Memory API (Third Preview)</a>
 */
@SuppressWarnings("OptionalGetWithoutIsPresent")
public class OpenSSL {
	// for selection
	public static final int EVP_PKEY_KEY_PARAMETERS = 0x84;
	public static final int EVP_PKEY_PUBLIC_KEY = 0x86;
	public static final int EVP_PKEY_KEYPAIR = 0x87;
	// for OSSL_PARAM.data_type
	public static final int OSSL_PARAM_INTEGER = 1; // native endian
	public static final int OSSL_PARAM_UNSIGNED_INTEGER = 2; // native endian
	public static final int OSSL_PARAM_REAL = 3; // native endian
	public static final int OSSL_PARAM_UTF8_STRING = 4;
	public static final int OSSL_PARAM_OCTET_STRING = 5;
	public static final int OSSL_PARAM_UTF8_PTR = 6;
	public static final int OSSL_PARAM_OCTET_PTR = 7;
	// evp methods
	public static final MethodHandle EVP_PKEY_Q_keygen;
	public static final MethodHandle EVP_PKEY_CTX_new;
	public static final MethodHandle EVP_PKEY_CTX_new_from_pkey;
	public static final MethodHandle EVP_PKEY_todata;
	public static final MethodHandle EVP_PKEY_fromdata_init;
	public static final MethodHandle EVP_PKEY_fromdata;
	public static final MethodHandle OSSL_PARAM_free;
	public static final MethodHandle EVP_PKEY_encrypt_init;
	public static final MethodHandle EVP_PKEY_encrypt;
	public static final MethodHandle EVP_PKEY_decrypt_init;
	public static final MethodHandle EVP_PKEY_decrypt;
	public static final MethodHandle EVP_PKEY_derive_init;
	public static final MethodHandle EVP_PKEY_derive_set_peer;
	public static final MethodHandle EVP_PKEY_derive;
	public static final MethodHandle EVP_PKEY_CTX_free;
	public static final MethodHandle EVP_PKEY_free;
	// global const strings
	private static final MemorySegment str_RSA;
	private static final MemorySegment str_EC; // openssl ecparam -list_curves
	private static final MemorySegment str_secp256k1;
	@SuppressWarnings("unused")
	private static final StructLayout OSSL_PARAM;

	static {
		// from https://kb.firedaemon.com/support/solutions/articles/4000121705
		var libCrypto = SymbolLookup.libraryLookup("libcrypto-3-x64.dll", Arena.global());
		var linker = Linker.nativeLinker();

		// EVP_PKEY *EVP_PKEY_Q_keygen(OSSL_LIB_CTX *libctx, const char *propq, const char *type, ...);
		EVP_PKEY_Q_keygen = linker.downcallHandle(libCrypto.find("EVP_PKEY_Q_keygen").get(),
				FunctionDescriptor.of(JAVA_LONG, JAVA_LONG, JAVA_LONG, ADDRESS, JAVA_LONG));

		// EVP_PKEY_CTX *EVP_PKEY_CTX_new(EVP_PKEY *pkey, ENGINE *e);
		EVP_PKEY_CTX_new = linker.downcallHandle(libCrypto.find("EVP_PKEY_CTX_new").get(),
				FunctionDescriptor.of(JAVA_LONG, JAVA_LONG, JAVA_LONG));

		// EVP_PKEY_CTX *EVP_PKEY_CTX_new_from_pkey(OSSL_LIB_CTX *libctx, EVP_PKEY *pkey, const char *propquery);
		EVP_PKEY_CTX_new_from_pkey = linker.downcallHandle(libCrypto.find("EVP_PKEY_CTX_new_from_pkey").get(),
				FunctionDescriptor.of(JAVA_LONG, JAVA_LONG, JAVA_LONG, ADDRESS));

		// int EVP_PKEY_todata(const EVP_PKEY *pkey, int selection, OSSL_PARAM **params);
		EVP_PKEY_todata = linker.downcallHandle(libCrypto.find("EVP_PKEY_todata").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, JAVA_INT, ADDRESS));

		// int EVP_PKEY_fromdata_init(EVP_PKEY_CTX *ctx);
		EVP_PKEY_fromdata_init = linker.downcallHandle(libCrypto.find("EVP_PKEY_fromdata_init").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG));

		// int EVP_PKEY_fromdata(EVP_PKEY_CTX *ctx, EVP_PKEY **ppkey, int selection, OSSL_PARAM params[]);
		EVP_PKEY_fromdata = linker.downcallHandle(libCrypto.find("EVP_PKEY_fromdata").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, JAVA_INT, JAVA_LONG));

		// void OSSL_PARAM_free(OSSL_PARAM *p);
		OSSL_PARAM_free = linker.downcallHandle(libCrypto.find("OSSL_PARAM_free").get(),
				FunctionDescriptor.ofVoid(JAVA_LONG));

		// int EVP_PKEY_encrypt_init(EVP_PKEY_CTX *ctx);
		EVP_PKEY_encrypt_init = linker.downcallHandle(libCrypto.find("EVP_PKEY_encrypt_init").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG));

		// int EVP_PKEY_encrypt(EVP_PKEY_CTX *ctx, unsigned char *out, size_t *outlen, const unsigned char *in, size_t inlen);
		EVP_PKEY_encrypt = linker.downcallHandle(libCrypto.find("EVP_PKEY_encrypt").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, ADDRESS, ADDRESS, JAVA_LONG));

		// int EVP_PKEY_decrypt_init(EVP_PKEY_CTX *ctx);
		EVP_PKEY_decrypt_init = linker.downcallHandle(libCrypto.find("EVP_PKEY_decrypt_init").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG));

		// int EVP_PKEY_decrypt(EVP_PKEY_CTX *ctx, unsigned char *out, size_t *outlen, const unsigned char *in, size_t inlen);
		EVP_PKEY_decrypt = linker.downcallHandle(libCrypto.find("EVP_PKEY_decrypt").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, ADDRESS, ADDRESS, JAVA_LONG));

		// int EVP_PKEY_derive_init(EVP_PKEY_CTX *ctx);
		EVP_PKEY_derive_init = linker.downcallHandle(libCrypto.find("EVP_PKEY_derive_init").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG));

		// int EVP_PKEY_derive_set_peer(EVP_PKEY_CTX *ctx, EVP_PKEY *peer);
		EVP_PKEY_derive_set_peer = linker.downcallHandle(libCrypto.find("EVP_PKEY_derive_set_peer").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, JAVA_LONG));

		// int EVP_PKEY_derive(EVP_PKEY_CTX *ctx, unsigned char *key, size_t *keylen);
		EVP_PKEY_derive = linker.downcallHandle(libCrypto.find("EVP_PKEY_derive").get(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, ADDRESS));

		// void EVP_PKEY_CTX_free(EVP_PKEY_CTX *ctx);
		EVP_PKEY_CTX_free = linker.downcallHandle(libCrypto.find("EVP_PKEY_CTX_free").get(),
				FunctionDescriptor.ofVoid(JAVA_LONG));

		// void EVP_PKEY_free(EVP_PKEY *pkey);
		EVP_PKEY_free = linker.downcallHandle(libCrypto.find("EVP_PKEY_free").get(),
				FunctionDescriptor.ofVoid(JAVA_LONG));

		OSSL_PARAM = MemoryLayout.structLayout( // [28]
				ADDRESS.withName("key"), //[+00] const char *key; /* the name of the parameter */
				JAVA_LONG.withName("data_type"), //[+08] unsigned char data_type; /* declare what kind of content is in data */
				ADDRESS.withName("data"), //[+10] void *data; /* value being passed in or out */
				JAVA_LONG.withName("data_size"), //[+18] size_t data_size; /* data size */
				JAVA_LONG.withName("return_size") //[+20] size_t return_size; /* returned size */
		);

		var global = Arena.global();
		str_RSA = global.allocateFrom("RSA", StandardCharsets.UTF_8);
		str_EC = global.allocateFrom("EC", StandardCharsets.UTF_8);
		str_secp256k1 = global.allocateFrom("secp256k1", StandardCharsets.UTF_8);
	}

	public static void dumpOsslParams(long osslParams) {
		var ms = MemorySegment.ofAddress(osslParams);
		for (long offset = 0; ; offset += 0x28) {
			ms = ms.reinterpret(offset + 0x28);
			var keyMs = ms.get(ADDRESS, offset);
			if (keyMs.address() == 0)
				break;
			keyMs = keyMs.reinterpret(1024);
			var key = keyMs.getString(0, StandardCharsets.UTF_8);
			var dataType = ms.get(JAVA_BYTE, offset + 0x8);
			var dataSize = ms.get(JAVA_LONG, offset + 0x18);
			var returnSize = ms.get(JAVA_LONG, offset + 0x20);
			System.out.println("key='" + key
					+ "', dataType=" + dataType
					+ ", dataSize=" + dataSize
					+ ", returnSize=" + returnSize);
		}
	}

	public static void testRsa() throws Throwable {
		var t = System.nanoTime();
		var evp_pkey = (long)EVP_PKEY_Q_keygen.invoke(0, 0, str_RSA, 2048L);
		System.out.println("EVP_PKEY_Q_keygen: " + evp_pkey + ", " + (System.nanoTime() - t) + "ns");
		// var evp_pkey_ctx = (long)EVP_PKEY_CTX_new_from_pkey.invoke(0, evp_pkey, MemorySegment.NULL);
		var evp_pkey_ctx = (long)EVP_PKEY_CTX_new.invoke(evp_pkey, 0);
		System.out.println("EVP_PKEY_CTX_new: " + evp_pkey_ctx);

		try (var arena = Arena.ofConfined()) {
			var osslParamsMs = arena.allocate(ADDRESS);
			var r = EVP_PKEY_todata.invoke(evp_pkey, EVP_PKEY_KEY_PARAMETERS, osslParamsMs);
			System.out.println("EVP_PKEY_todata(EVP_PKEY_KEY_PARAMETERS): " + r);
			var osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_todata.invoke(evp_pkey, EVP_PKEY_KEYPAIR, osslParamsMs);
			System.out.println("EVP_PKEY_todata(EVP_PKEY_KEYPAIR): " + r);
			osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_todata.invoke(evp_pkey, EVP_PKEY_PUBLIC_KEY, osslParamsMs);
			System.out.println("EVP_PKEY_todata(EVP_PKEY_PUBLIC_KEY): " + r);
			osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_encrypt_init.invoke(evp_pkey_ctx);
			System.out.println("EVP_PKEY_encrypt_init: " + r);
			var srcBuf = arena.allocateFrom("123", StandardCharsets.UTF_8);
			var srcBufSize = srcBuf.byteSize();
			var encBuf = arena.allocate(1024, 1);
			var encSizeBuf = arena.allocate(8, 8);
			System.out.println("EVP_PKEY_encrypt: size=" + srcBufSize
					+ ", encBuf.size=" + encBuf.byteSize()
					+ ", encSizeBuf.size=" + encSizeBuf.byteSize());
			for (int i = 0; i < 1000; i++) {
				encSizeBuf.set(JAVA_LONG, 0, 1024);
				t = System.nanoTime();
				r = EVP_PKEY_encrypt.invoke(evp_pkey_ctx, encBuf, encSizeBuf, srcBuf, srcBufSize);
			}
			t = System.nanoTime() - t;
			var encSize = encSizeBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_encrypt: " + r + ", encSize=" + encSize + ", " + t + "ns");

			r = EVP_PKEY_decrypt_init.invoke(evp_pkey_ctx);
			System.out.println("EVP_PKEY_decrypt_init: " + r);
			var dstBuf = arena.allocate(1024, 1);
			var dstSizeBuf = arena.allocate(8, 8);
			for (int i = 0; i < 1000; i++) {
				dstSizeBuf.set(JAVA_LONG, 0, 1024);
				t = System.nanoTime();
				r = EVP_PKEY_decrypt.invoke(evp_pkey_ctx, dstBuf, dstSizeBuf, encBuf, encSize);
			}
			t = System.nanoTime() - t;
			System.out.println("EVP_PKEY_decrypt: " + r
					+ ", dstSize=" + dstSizeBuf.get(JAVA_LONG, 0) + ", " + t + "ns");
			System.out.println("EVP_PKEY_decrypt: "
					+ dstBuf.get(JAVA_BYTE, 0) + ", "
					+ dstBuf.get(JAVA_BYTE, 1) + ", "
					+ dstBuf.get(JAVA_BYTE, 2) + ", "
					+ dstBuf.get(JAVA_BYTE, 3));
		} finally {
			if (evp_pkey_ctx != 0)
				EVP_PKEY_CTX_free.invoke(evp_pkey_ctx);
			if (evp_pkey != 0)
				EVP_PKEY_free.invoke(evp_pkey);
			System.out.println("END");
		}
	}

	public static void testEc() throws Throwable {
		var t = System.nanoTime();
		var evp_pkey0 = (long)EVP_PKEY_Q_keygen.invoke(0, 0, str_EC, str_secp256k1.address());
		var evp_pkey1 = (long)EVP_PKEY_Q_keygen.invoke(0, 0, str_EC, str_secp256k1.address());
		long evp_pubkey0 = 0;
		long evp_pubkey1 = 0;
		System.out.println("EVP_PKEY_Q_keygen0: " + evp_pkey0 + ", " + (System.nanoTime() - t) + "ns");
		System.out.println("EVP_PKEY_Q_keygen1: " + evp_pkey1 + ", " + (System.nanoTime() - t) + "ns");
		var evp_pkey_ctx0 = (long)EVP_PKEY_CTX_new.invoke(evp_pkey0, 0);
		var evp_pkey_ctx1 = (long)EVP_PKEY_CTX_new.invoke(evp_pkey1, 0);
		System.out.println("EVP_PKEY_CTX_new0: " + evp_pkey_ctx0);
		System.out.println("EVP_PKEY_CTX_new1: " + evp_pkey_ctx1);

		try (var arena = Arena.ofConfined()) {
			var osslParamsMs = arena.allocate(ADDRESS);
			var r = EVP_PKEY_todata.invoke(evp_pkey0, EVP_PKEY_KEY_PARAMETERS, osslParamsMs);
			System.out.println("EVP_PKEY_todata(EVP_PKEY_KEY_PARAMETERS): " + r);
			var osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_todata.invoke(evp_pkey0, EVP_PKEY_KEYPAIR, osslParamsMs);
			System.out.println("EVP_PKEY_todata(EVP_PKEY_KEYPAIR): " + r);
			osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_todata.invoke(evp_pkey0, EVP_PKEY_PUBLIC_KEY, osslParamsMs);
			System.out.println("EVP_PKEY_todata0(EVP_PKEY_PUBLIC_KEY): " + r);
			osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);

			r = EVP_PKEY_fromdata_init.invoke(evp_pkey_ctx0);
			System.out.println("EVP_PKEY_fromdata_init0: " + r);
			var pubKeyBuf = arena.allocate(8, 8);
			pubKeyBuf.set(JAVA_LONG, 0, 0);
			r = EVP_PKEY_fromdata.invoke(evp_pkey_ctx0, pubKeyBuf, EVP_PKEY_PUBLIC_KEY, osslParams);
			evp_pubkey0 = pubKeyBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_fromdata0: " + r + ", pubKey=" + evp_pubkey0);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_todata.invoke(evp_pkey1, EVP_PKEY_PUBLIC_KEY, osslParamsMs);
			System.out.println("EVP_PKEY_todata1(EVP_PKEY_PUBLIC_KEY): " + r);
			osslParams = osslParamsMs.get(JAVA_LONG, 0);
			dumpOsslParams(osslParams);

			r = EVP_PKEY_fromdata_init.invoke(evp_pkey_ctx1);
			System.out.println("EVP_PKEY_fromdata_init1: " + r);
			pubKeyBuf.set(JAVA_LONG, 0, 0);
			r = EVP_PKEY_fromdata.invoke(evp_pkey_ctx1, pubKeyBuf, EVP_PKEY_PUBLIC_KEY, osslParams);
			evp_pubkey1 = pubKeyBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_fromdata1: " + r + ", pubKey=" + evp_pubkey1);
			OSSL_PARAM_free.invoke(osslParams);

			r = EVP_PKEY_derive_init.invoke(evp_pkey_ctx0);
			System.out.println("EVP_PKEY_encrypt_init0: " + r);
			r = EVP_PKEY_derive_init.invoke(evp_pkey_ctx1);
			System.out.println("EVP_PKEY_encrypt_init1: " + r);

			r = EVP_PKEY_derive_set_peer.invoke(evp_pkey_ctx0, evp_pubkey1);
			System.out.println("EVP_PKEY_derive_set_peer0: " + r);
			r = EVP_PKEY_derive_set_peer.invoke(evp_pkey_ctx1, evp_pubkey0);
			System.out.println("EVP_PKEY_derive_set_peer1: " + r);

			var secretSizeBuf = arena.allocate(8, 8);
			r = EVP_PKEY_derive.invoke(evp_pkey_ctx0, MemorySegment.NULL, secretSizeBuf);
			var secretSize0 = secretSizeBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_derive0: " + r + ", secretSize=" + secretSize0);
			var secretBuf0 = arena.allocate(secretSize0, 1);
			r = EVP_PKEY_derive.invoke(evp_pkey_ctx0, secretBuf0, secretSizeBuf);
			secretSize0 = secretSizeBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_derive0: " + r + ", secretSize=" + secretSize0);

			r = EVP_PKEY_derive.invoke(evp_pkey_ctx0, MemorySegment.NULL, secretSizeBuf);
			var secretSize1 = secretSizeBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_derive1: " + r + ", secretSize=" + secretSize1);
			var secretBuf1 = arena.allocate(secretSize1, 1);
			r = EVP_PKEY_derive.invoke(evp_pkey_ctx0, secretBuf1, secretSizeBuf);
			secretSize1 = secretSizeBuf.get(JAVA_LONG, 0);
			System.out.println("EVP_PKEY_derive1: " + r + ", secretSize=" + secretSize1);

			if (secretSize0 == secretSize1) {
				for (int i = 0; ; ) {
					if (secretBuf0.get(JAVA_BYTE, i) != secretBuf1.get(JAVA_BYTE, i))
						break;
					if (++i == secretSize0) {
						System.out.println("secret check OK!");
						break;
					}
				}
			}
		} finally {
			if (evp_pkey_ctx0 != 0)
				EVP_PKEY_CTX_free.invoke(evp_pkey_ctx0);
			if (evp_pkey_ctx1 != 0)
				EVP_PKEY_CTX_free.invoke(evp_pkey_ctx1);
			if (evp_pubkey0 != 0)
				EVP_PKEY_free.invoke(evp_pubkey0);
			if (evp_pubkey1 != 0)
				EVP_PKEY_free.invoke(evp_pubkey1);
			if (evp_pkey0 != 0)
				EVP_PKEY_free.invoke(evp_pkey0);
			if (evp_pkey1 != 0)
				EVP_PKEY_free.invoke(evp_pkey1);
			System.out.println("END");
		}
	}

	public static void main(String[] args) throws Throwable {
		// testRsa();
		testEc();
	}
}
